Upload predict.py

predict.py

@@ -0,0 +1,104 @@
+# Prediction interface for Cog ⚙️
+# https://github.com/replicate/cog/blob/main/docs/python.md
+
+from cog import BasePredictor, Input, Path
+import tempfile
+import os, glob
+import numpy as np
+import cv2
+from PIL import Image
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from models import model, basic
+from utils import util
+
+class Predictor(BasePredictor):
+    def setup(self):
+        seed = 130
+        np.random.seed(seed)
+        torch.manual_seed(seed)
+        torch.cuda.manual_seed(seed)
+        #print('--------------', torch.cuda.is_available())
+        """Load the model into memory to make running multiple predictions efficient"""
+        self.colorizer = model.AnchorColorProb(inChannel=1, outChannel=313, enhanced=True)
+        self.colorizer = self.colorizer.cuda()
+        checkpt_path = "./checkpoints/disco-beta.pth.rar"
+        assert os.path.exists(checkpt_path)
+        data_dict = torch.load(checkpt_path, map_location=torch.device('cpu'))
+        self.colorizer.load_state_dict(data_dict['state_dict'])
+        self.colorizer.eval()
+        self.color_class = basic.ColorLabel(lambda_=0.5, device='cuda')
+    
+    def resize_ab2l(self, gray_img, lab_imgs):
+        H, W = gray_img.shape[:2]
+        reszied_ab = cv2.resize(lab_imgs[:,:,1:], (W,H), interpolation=cv2.INTER_LINEAR)
+        return np.concatenate((gray_img, reszied_ab), axis=2)
+    
+    def predict(
+        self,
+        image: Path = Input(description="input image. Output will be one or multiple colorized images."),
+        n_anchors: int = Input(
+            description="number of color anchors", ge=3, le=14, default=8
+        ),
+        multi_result: bool = Input(
+            description="to generate diverse results", default=False
+        ),
+        vis_anchors: bool = Input(
+            description="to visualize the anchor locations", default=False
+        )
+    ) -> Path:
+        """Run a single prediction on the model"""
+        bgr_img = cv2.imread(str(image), cv2.IMREAD_COLOR)
+        rgb_img = cv2.cvtColor(bgr_img, cv2.COLOR_BGR2RGB)
+        rgb_img = np.array(rgb_img / 255., np.float32)
+        lab_img = cv2.cvtColor(rgb_img, cv2.COLOR_RGB2LAB)
+        org_grays = (lab_img[:,:,[0]]-50.) / 50.
+        lab_img = cv2.resize(lab_img, (256,256), interpolation=cv2.INTER_LINEAR)
+        
+        lab_img = torch.from_numpy(lab_img.transpose((2, 0, 1)))
+        gray_img = (lab_img[0:1,:,:]-50.) / 50.
+        ab_chans = lab_img[1:3,:,:] / 110.
+        input_grays = gray_img.unsqueeze(0)
+        input_colors = ab_chans.unsqueeze(0)
+        input_grays = input_grays.cuda(non_blocking=True)
+        input_colors = input_colors.cuda(non_blocking=True)
+                
+        sampled_T = 2 if multi_result else 0
+        pal_logit, ref_logit, enhanced_ab, affinity_map, spix_colors, hint_mask = self.colorizer(input_grays, \
+                                                            input_colors, n_anchors, True, sampled_T)
+        pred_probs = pal_logit
+        guided_colors = self.color_class.decode_ind2ab(ref_logit, T=0)
+        sp_size = 16
+        guided_colors = basic.upfeat(guided_colors, affinity_map, sp_size, sp_size)
+        res_list = []
+        if multi_result:
+            for no in range(3):
+                pred_labs = torch.cat((input_grays,enhanced_ab[no:no+1,:,:,:]), dim=1)
+                lab_imgs = basic.tensor2array(pred_labs).squeeze(axis=0)
+                lab_imgs = self.resize_ab2l(org_grays, lab_imgs)
+                #util.save_normLabs_from_batch(lab_imgs, save_dir, [file_name], -1, suffix='c%d'%no)
+                res_list.append(lab_imgs)
+        else:
+            pred_labs = torch.cat((input_grays,enhanced_ab), dim=1)
+            lab_imgs = basic.tensor2array(pred_labs).squeeze(axis=0)
+            lab_imgs = self.resize_ab2l(org_grays, lab_imgs)
+            #util.save_normLabs_from_batch(lab_imgs, save_dir, [file_name], -1)#, suffix='enhanced')
+            res_list.append(lab_imgs)
+        
+        if vis_anchors:
+            ## visualize anchor locations
+            anchor_masks = basic.upfeat(hint_mask, affinity_map, sp_size, sp_size)
+            marked_labs = basic.mark_color_hints(input_grays, enhanced_ab, anchor_masks, base_ABs=enhanced_ab)
+            hint_imgs = basic.tensor2array(marked_labs).squeeze(axis=0)
+            hint_imgs = self.resize_ab2l(org_grays, hint_imgs)
+            #util.save_normLabs_from_batch(hint_imgs, save_dir, [file_name], -1, suffix='anchors')
+            res_list.append(hint_imgs)
+        
+        output = cv2.vconcat(res_list)
+        output[:,:,0] = output[:,:,0] * 50.0 + 50.0
+        output[:,:,1:3] = output[:,:,1:3] * 110.0
+        rgb_output = cv2.cvtColor(output[:,:,:], cv2.COLOR_LAB2BGR)      
+        out_path = Path(tempfile.mkdtemp()) / "out.png"
+        cv2.imwrite(str(out_path), (rgb_output*255.0).astype(np.uint8))
+        return out_path