pruning

sam_utils.py

@@ -1,14 +1,10 @@
 import os
 import numpy as np
 import torch
-# import matplotlib.pyplot as plt
-import cv2
 from PIL import Image
-# from PIL import Image
 import time
-from utils import find_image_file
 
-from segment_anything import sam_model_registry, SamPredictor, SamAutomaticMaskGenerator
+from segment_anything import sam_model_registry, SamPredictor
 
 def sam_init(device_id=0):
     import inspect
@@ -22,60 +18,11 @@ def sam_init(device_id=0):
     sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)
     sam.to(device=device)
     predictor = SamPredictor(sam)
-    # mask_generator = SamAutomaticMaskGenerator(sam)
     return predictor
 
-def sam_out(predictor, shape_dir):
-    image_path = os.path.join(shape_dir, find_image_file(shape_dir))
-    save_path = os.path.join(shape_dir, "image_sam.png")
-    bbox_path = os.path.join(shape_dir, "bbox.txt")
-    bbox = np.loadtxt(bbox_path, delimiter=',')
-    image = cv2.imread(image_path)
-    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
-
-    start_time = time.time()
-    predictor.set_image(image)
-
-    h, w, _ = image.shape
-    input_point = np.array([[h//2, w//2]])
-    input_label = np.array([1])
-
-    masks, scores, logits = predictor.predict(
-        point_coords=input_point,
-        point_labels=input_label,
-        multimask_output=True,
-    )
-
-    masks_bbox, scores_bbox, logits_bbox = predictor.predict(
-        box=bbox,
-        multimask_output=True
-    )
-
-    print(f"SAM Time: {time.time() - start_time:.3f}s")
-    opt_idx = np.argmax(scores)
-    mask = masks[opt_idx]
-    out_image = np.zeros((image.shape[0], image.shape[1], 4), dtype=np.uint8)
-    out_image[:, :, :3] = image
-    out_image_bbox = out_image.copy()
-    out_image[:, :, 3] = mask.astype(np.uint8) * 255
-    out_image_bbox[:, :, 3] = masks_bbox[-1].astype(np.uint8) * 255 # np.argmax(scores_bbox)
-    cv2.imwrite(save_path, cv2.cvtColor(out_image_bbox, cv2.COLOR_RGBA2BGRA))
-
-
-def convert_from_cv2_to_image(img: np.ndarray) -> Image:
-    return Image.fromarray(img)
-    # return Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGRA2RGBA))
-
-def convert_from_image_to_cv2(img: Image) -> np.ndarray:
-    return np.asarray(img)
-    # return cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)
-
 def sam_out_nosave(predictor, input_image, *bbox_sliders):
-    # save_path = os.path.join(shape_dir, "image_sam.png")
-    # bbox_path = os.path.join(shape_dir, "bbox.txt")
-    # bbox = np.loadtxt(bbox_path, delimiter=',')
     bbox = np.array(bbox_sliders)
-    image = convert_from_image_to_cv2(input_image)
+    image = np.asarray(input_image)
 
     start_time = time.time()
     predictor.set_image(image)
@@ -104,5 +51,4 @@ def sam_out_nosave(predictor, input_image, *bbox_sliders):
     out_image[:, :, 3] = mask.astype(np.uint8) * 255
     out_image_bbox[:, :, 3] = masks_bbox[-1].astype(np.uint8) * 255 # np.argmax(scores_bbox)
     torch.cuda.empty_cache()
-    return Image.fromarray(out_image_bbox, mode='RGBA') 
-    cv2.imwrite(save_path, cv2.cvtColor(out_image_bbox, cv2.COLOR_RGBA2BGRA))
+    return Image.fromarray(out_image_bbox, mode='RGBA') 

zero123_utils.py

@@ -76,7 +76,7 @@ def sample_model_batch(model, sampler, input_im, xs, ys, n_samples=4, precision=
             cond = {}
             cond['c_crossattn'] = [c]
             # c_concat = model.encode_first_stage((input_im.to(c.device))).mode().detach()
-            cond['c_concat'] = [model.encode_first_stage((input_im.to(c.device))).mode().detach()
+            cond['c_concat'] = [model.encode_first_stage(input_im).mode().detach()
                                 .repeat(n_samples, 1, 1, 1)]
             if scale != 1.0:
                 uc = {}
@@ -99,7 +99,8 @@ def sample_model_batch(model, sampler, input_im, xs, ys, n_samples=4, precision=
             # samples_ddim = torch.nn.functional.interpolate(samples_ddim, 64, mode='nearest', antialias=False)
             x_samples_ddim = model.decode_first_stage(samples_ddim)
             ret_imgs = torch.clamp((x_samples_ddim + 1.0) / 2.0, min=0.0, max=1.0).cpu()
-            del cond, c, x_samples_ddim, samples_ddim, uc
+            del cond, c, x_samples_ddim, samples_ddim, uc, input_im
+            torch.cuda.empty_cache()
             return ret_imgs
         
 
@@ -126,6 +127,7 @@ def predict_stage1(model, sampler, input_img_path, save_path_8, adjust_set=[], d
     del input_im
     torch.cuda.empty_cache()
 
+@torch.no_grad()
 def predict_stage1_gradio(model, raw_im, save_path = "", adjust_set=[], device="cuda", ddim_steps=75, scale=3.0):
     # raw_im = raw_im.resize([256, 256], Image.LANCZOS)
     # input_im_init = preprocess_image(models, raw_im, preprocess=False)
@@ -157,7 +159,6 @@ def predict_stage1_gradio(model, raw_im, save_path = "", adjust_set=[], device="
             out_image.save(os.path.join(save_path, '%d.png'%(stage1_idx)))
         sample_idx += 1
     del x_samples_ddims_8
-    del input_im
     del sampler
     torch.cuda.empty_cache()
     return ret_imgs
@@ -188,7 +189,6 @@ def infer_stage_2(model, save_path_stage1, save_path_stage2, delta_x_2, delta_y_
             x_sample_stage2 = 255.0 * rearrange(x_samples_ddims_stage2[stage2_idx].numpy(), 'c h w -> h w c')
             Image.fromarray(x_sample_stage2.astype(np.uint8)).save(os.path.join(save_path_stage2, '%d_%d.png'%(stage1_idx, stage2_idx)))
         del input_im
-        del sampler
         del x_samples_ddims_stage2
         torch.cuda.empty_cache()