flix-clip2 0.0.1__py3-none-any.whl

flix_clip2/__init__.py

@@ -0,0 +1,4 @@
+import os 
+import sys
+path = os.path.dirname(os.path.realpath(__file__))
+sys.path.append(path)

flix_clip2/api.py

@@ -0,0 +1,214 @@
+import os
+import json
+import cv2
+import numpy as np
+from time import time
+from termcolor import colored
+import shutil
+
+from dis_flix.api import Bolt as Bolt_dis
+from matteformer.api import Bolt as Bolt_matteformer
+from objdetect_pytorch import Bolt as Bolt_obj
+
+UPDATE = "25 Oct 2023"
+
+class Bolt:
+    def __init__(self, init_dict_main):
+        self.flag_detect = init_dict_main['detect']
+        self.flag_fine = init_dict_main['fine_segmentation']
+        self.single_object = True #init_dict_main['single_object']
+        self.gpu_id = init_dict_main['gpu_id']
+        self.temp_dir = init_dict_main['temp_dir']
+        self.house_keeping = init_dict_main['house_keeping']
+
+        os.makedirs(self.temp_dir, exist_ok=True)
+
+        # ------------ Intro ------------
+        bolt_names = ["DIS"]
+        if self.flag_fine:
+            bolt_names.append("Matteformer")
+        if self.flag_detect:
+            bolt_names.append("Object Detection")
+        msg = "Bolt (%s)"%(" + ".join(bolt_names))
+        msg = msg + " | Update: %s"%(UPDATE)
+        print_special(msg,color="yellow",line_char="=",n=100)
+
+        # ------------ Bolt DIS------------
+        init_dict = {'model_path':init_dict_main['coarse_model_path'],'input_size':init_dict_main['dis_input_size'],'gpu_id':self.gpu_id}
+        self.bolt_dis = Bolt_dis(init_dict)
+        print(colored('DIS Bolt Initiallized','yellow'))
+        print(colored("-"*20, "yellow"))
+
+        # ------------ Bolt Matteformer------------
+        if self.flag_fine:
+            # load config
+            # config_path = init_dict_main['matte_config_path']
+            # with open(config_path, "r") as f:
+            #     config = json.load(f)
+            # Initialize Bolt
+            config=init_dict_main['matte_config_dict']
+            config["gpu_id"] = self.gpu_id
+            self.bolt_matte = Bolt_matteformer(config)
+            print(colored('Mattefomer Bolt Initiallized (new)','yellow'))
+            print(colored("-"*20, "yellow"))
+
+        # ------------ Bolt Object Detection------------
+        if self.flag_detect:
+            gpu_id = self.gpu_id
+            assert os.path.isfile(init_dict_main['object_detect_model_path']), "Object detection model not found: %s"%(init_dict_main['object_detect_model_path'])
+            checkpoint_dict = {
+                                "obj_model": init_dict_main['object_detect_model_path'],
+                                "category_count": 2
+                            }
+
+            self.obj = Bolt_obj(checkpoint_dict, gpu_id)
+            print(colored('Object Detection Bolt Initiallized','yellow'))
+            print(colored("-"*20, "yellow"))
+
+
+    def run(self, input_dict):
+        try:
+            tic_all = time()
+            image_path = input_dict['image_path']
+            temp_path = os.path.join(self.temp_dir, str(time()))
+            os.makedirs(temp_path, exist_ok=True)
+            # ------------ Object Detection ------------
+            if self.flag_detect:
+                tic = time()
+                detect_path = os.path.join(temp_path, "detect.jpg")
+                input_dict_obj = {
+                                "filepath": image_path,
+                                "setid": 'sku_id',
+                                "out_dir": ""
+                            }
+                flag, crop_img_list, detection_box_list = self.obj.run(input_dict_obj)
+                if flag:
+                    if self.single_object:
+                        cv2.imwrite(detect_path, crop_img_list[0])
+                else:
+                    err_msg = "Detection FAILED!!!!"
+                    raise Exception(err_msg)
+                toc = time()
+                print(colored('Time (Detect): %0.2f'%(toc-tic),'yellow'))
+                print(colored("-"*20, "yellow"))
+            
+            # ------------ DIS ------------
+            tic = time()
+            if self.flag_detect:
+                image_path = detect_path
+                result_detect_path = os.path.join(temp_path, "result.png")
+            else:
+                result_detect_path = input_dict['result_path'] # final result
+            
+            if self.flag_fine:
+                segment_path = os.path.join(temp_path, "coarse.png")
+                fine_path = result_detect_path
+            else:
+                segment_path = result_detect_path
+                
+            input_dict_dis = {'image_path': image_path,'result_path': segment_path}
+            os.makedirs(os.path.dirname(input_dict_dis["result_path"]), exist_ok=True)
+            out_dict_dis, status = self.bolt_dis.run(input_dict_dis)
+            toc = time()
+            print(colored('Time (Coarse): %0.2f'%(toc-tic),'yellow'))
+            if status:
+                if self.flag_fine:
+                    assert os.path.isfile(out_dict_dis['result_path']), "Coarse segmentation not found: %s"%(out_dict_dis['result_path'])
+                    tic_matte = time()
+                    # ------------ Infer (fine)------------
+                    # Create input dictionary
+                    input_dict_fine = {"image_path": image_path,
+                                "segment_path": segment_path,
+                                "output_path": fine_path}
+                    os.makedirs(os.path.dirname(input_dict_fine["output_path"]), exist_ok=True)
+                    # Run Bolt
+                    self.bolt_matte.run(input_dict_fine)
+                        
+                    toc_matte = time()
+                    print(colored('Time (Matte): %0.2f'%(toc_matte-tic_matte),'yellow'))
+            else:
+                err_msg = out_dict_dis['error_status']
+                raise Exception(err_msg)
+            
+            # ------------ Construct Original Mask ------------
+            if self.flag_detect:
+                org = cv2.imread(input_dict['image_path'], 0)
+                clip_detect = cv2.imread(result_detect_path, 0)
+                canvas = self.unpad_result(org, clip_detect, detection_box_list)
+                os.makedirs(os.path.dirname(input_dict['result_path']), exist_ok=True)
+                cv2.imwrite(input_dict['result_path'], canvas)
+                
+                # ------------save coarse ------------
+                if "coarse_path" in input_dict:
+                    coarse_path = input_dict['coarse_path']
+                    clip_detect = cv2.imread(segment_path, 0)
+                    canvas = self.unpad_result(org, clip_detect, detection_box_list)
+                    os.makedirs(os.path.dirname(coarse_path), exist_ok=True)
+                    cv2.imwrite(coarse_path, canvas)
+                    print(colored("Coarse saved: %s"%(coarse_path),"cyan","on_yellow"))
+                # else: result is saved as required already
+            else:
+                # ------------save coarse ------------
+                if "coarse_path" in input_dict:
+                    coarse_path = input_dict['coarse_path']
+                    os.makedirs(os.path.dirname(coarse_path), exist_ok=True)
+                    shutil.copyfile(segment_path, coarse_path)
+
+            # ------------ House Keeping ------------
+            if self.house_keeping:
+                shutil.rmtree(temp_path)
+                print(colored("House keeping done: %s"%temp_path,"blue"))
+
+            status = True
+            toc_all = time()
+            print(colored('Time (Total): %0.2f'%(toc_all-tic_all),'cyan'))
+            print(colored("-"*100, "cyan"))
+        except Exception as e:
+            err_msg = str(e)
+            status = False
+            print(colored(err_msg, "red"))
+            import traceback
+            traceback.print_exc()
+        
+        return status
+
+
+    def unpad_result(self, org, clip_detect, detection_box_list):
+        canv = np.zeros_like(org)
+        left, top, right, bottom = detection_box_list[0]['xmin'], detection_box_list[0]['ymin'], detection_box_list[0]['xmax'], detection_box_list[0]['ymax']
+        canv[top:bottom, left:right] = clip_detect
+        return canv
+
+def print_special(msg,color="yellow",line_char="-",n=100):
+    schar = " "
+    nfill = (n-len(msg)+2)//2 # number of fillers
+    msg = schar*nfill+" "+msg+" "+schar*nfill
+    print(colored(line_char*n,color=color))
+    print(colored(msg,color=color))
+    print(colored(line_char*n,color=color))
+
+
+if __name__ == "__main__":
+    init_dict_main = {}
+    init_dict_main['detect'] = True
+    init_dict_main['object_detect_model_path'] = '/datarive1/projects/python/beer/DeploymentMaster/flix_clip/packages/object_detect_package/step_075000.pth'
+    init_dict_main['coarse_model_path'] = '/datarive1/projects/python/DIS/saved_models/IS-Net/test_v3.pth'
+    init_dict_main['dis_input_size'] = [1024,1024]
+    init_dict_main['gpu_id'] = '0'
+    init_dict_main['fine_segmentation'] = True
+    init_dict_main['matte_config_path'] = '/datarive1/projects/python/beer/DeploymentMaster/flix_clip/packages/matteformer_package/infer_config.json'
+    init_dict_main['temp_dir'] = '/datarive1/projects/python/beer/DeploymentMaster/flix_clip/temp'
+    init_dict_main['house_keeping'] = True
+
+    # ------------ Bolt ------------
+    bolt = Bolt(init_dict_main)
+    # ------------ Infer ------------
+    input_dict = {}
+    input_dict['image_path'] = 'sample/images/356957_1_edit.jpg'
+    input_dict['result_path'] = 'sample/results/356957_1_edit.png'
+    # optional (only with fine_segmentation:true)
+    input_dict['coarse_path'] = 'sample/coarses/356957_1_edit.png'
+
+    flag = bolt.run(input_dict)
+    
+    print(flag)