jonstreeterandClaude
36860b09dd
This fixes the ModuleNotFoundError issues that occurred during clean installations. Changes: - Embedded colormnet_source dependencies directly into repository as: - colormnet_model/ (model architecture) - colormnet_inference/ (inference pipeline) - colormnet_util/ (utility functions) - colormnet_dataset/ (data transforms and utilities) - Updated all internal imports to use new embedded paths - Removed external colormnet_source path dependency from __init__.py - Added graceful fallback for spatial_correlation_sampler CUDA extension - Fixed shape calculation bugs in PyTorch fallback implementation - Added missing dependencies (hickle, scikit-image) to auto-install The repository is now self-contained and can be installed without requiring the separate colormnet_source directory. Tested and verified: - Node imports work correctly - Colorization pipeline runs successfully - Model loads and processes video frames Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <noreply@anthropic.com>
92 lines
2.9 KiB
Python
92 lines
2.9 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
|
|
from collections import defaultdict
|
|
|
|
|
|
def dice_loss(input_mask, cls_gt):
|
|
num_objects = input_mask.shape[1]
|
|
losses = []
|
|
for i in range(num_objects):
|
|
mask = input_mask[:,i].flatten(start_dim=1)
|
|
# background not in mask, so we add one to cls_gt
|
|
gt = (cls_gt==(i+1)).float().flatten(start_dim=1)
|
|
numerator = 2 * (mask * gt).sum(-1)
|
|
denominator = mask.sum(-1) + gt.sum(-1)
|
|
loss = 1 - (numerator + 1) / (denominator + 1)
|
|
losses.append(loss)
|
|
return torch.cat(losses).mean()
|
|
|
|
def l1_loss(input, target):
|
|
out = torch.abs(input - target)
|
|
return out.mean()
|
|
|
|
|
|
# https://stackoverflow.com/questions/63735255/how-do-i-compute-bootstrapped-cross-entropy-loss-in-pytorch
|
|
class BootstrappedCE(nn.Module):
|
|
def __init__(self, start_warm, end_warm, top_p=0.15):
|
|
super().__init__()
|
|
|
|
self.start_warm = start_warm
|
|
self.end_warm = end_warm
|
|
self.top_p = top_p
|
|
|
|
def forward(self, input, target, it):
|
|
if it < self.start_warm:
|
|
|
|
return F.cross_entropy(input, target), 1.0
|
|
|
|
raw_loss = F.cross_entropy(input, target, reduction='none').view(-1)
|
|
num_pixels = raw_loss.numel()
|
|
|
|
if it > self.end_warm:
|
|
this_p = self.top_p
|
|
else:
|
|
this_p = self.top_p + (1-self.top_p)*((self.end_warm-it)/(self.end_warm-self.start_warm))
|
|
loss, _ = torch.topk(raw_loss, int(num_pixels * this_p), sorted=False)
|
|
return loss.mean(), this_p
|
|
|
|
|
|
class LossComputer:
|
|
def __init__(self, config):
|
|
super().__init__()
|
|
self.config = config
|
|
self.bce = BootstrappedCE(config['start_warm'], config['end_warm'])
|
|
|
|
def compute(self, data, num_objects, it):
|
|
losses = defaultdict(int)
|
|
|
|
b, t = data['rgb'].shape[:2]
|
|
|
|
losses['total_loss'] = 0
|
|
for ti in range(1, t):
|
|
for bi in range(b):
|
|
loss, p = self.bce(data[f'logits_{ti}'][bi:bi+1, :num_objects[bi]+1], data['cls_gt'][bi:bi+1,ti,0], it)
|
|
losses['p'] += p / b / (t-1)
|
|
losses[f'ce_loss_{ti}'] += loss / b
|
|
|
|
losses['total_loss'] += losses['ce_loss_%d'%ti]
|
|
losses[f'dice_loss_{ti}'] = dice_loss(data[f'masks_{ti}'], data['cls_gt'][:,ti,0])
|
|
losses['total_loss'] += losses[f'dice_loss_{ti}']
|
|
|
|
return losses
|
|
|
|
|
|
def compute_l1loss(self, data, num_objects, it):
|
|
losses = defaultdict(int)
|
|
|
|
b, t = data['rgb'].shape[:2]
|
|
|
|
losses['total_loss'] = 0
|
|
for ti in range(1, t):
|
|
for bi in range(b):
|
|
losses['p'] = 0
|
|
losses[f'ce_loss_{ti}'] = 0
|
|
|
|
losses['total_loss'] += losses['ce_loss_%d'%ti]
|
|
losses[f'dice_loss_{ti}'] = l1_loss(data[f'masks_{ti}'], data['cls_gt'][:,ti])
|
|
losses['total_loss'] += losses[f'dice_loss_{ti}']
|
|
|
|
return losses
|