import torch import torchvision.transforms as T from torchvision.models.detection import maskrcnn_resnet50_fpn from PIL import Image import matplotlib.pyplot as plt import numpy as np import uuid import os import cv2 import json input_images_dir = 'data/input_images/' segmented_objects_dir = 'data/segmented_objects/' os.makedirs(input_images_dir, exist_ok=True) os.makedirs(segmented_objects_dir, exist_ok=True) #Loading the model def load_model(): model = maskrcnn_resnet50_fpn(pretrained=True) # Using a different backbone #model = maskrcnn_resnet50_fpn(pretrained=False, pretrained_backbone=False, backbone_name='resnext50_32x4d') model.eval() """ We have set this to evaluation mode, because we have loaded a pretrained model so we must deactivate dropout layers and other training-specific behaviors. """ return model model = load_model() #model initialization def transform_image(image): transform = T.Compose([ T.Resize((256, 256)), # Resize to match model input T.ToTensor(), # Convert to torch tensor T.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) # Normalize ]) return transform(image).unsqueeze(0) # Add batch dimension to get [1,C,H,W] #C is channels, RGB has 3, greyscale has 1 # # Test image transformation # image_path = "D:\multiobject.jpeg" # Replace with the path to your image # image_tensor = transform_image(image_path) def run_inference(model,image_tensor): with torch.no_grad(): outputs = model(image_tensor) return outputs def extract_object(image, mask): img_np = np.array(image) # Resize mask to match image dimensions mask_resized = cv2.resize(mask, (img_np.shape[1], img_np.shape[0]), interpolation=cv2.INTER_NEAREST) # Create an empty image with the same dimensions as the original image object_img = np.zeros_like(img_np) # Apply the mask to the image for c in range(3): # Assuming image has 3 channels (RGB) object_img[:, :, c] = img_np[:, :, c] * mask_resized return Image.fromarray(object_img) # def extract_object(image, mask): # object_img = Image.fromarray((np.array(image) * mask[:, :, None]).astype(np.uint8)) # return object_img # Save the input image def save_input_image(image, master_id): input_image_path = os.path.join(input_images_dir, f'{master_id}.png') image.save(input_image_path) return input_image_path # Save the extracted objects and their metadata def save_objects_and_metadata(extracted_objects, master_id): object_metadata = [] for i, obj_img in enumerate(extracted_objects): object_id = str(uuid.uuid4()) object_image_path = os.path.join(segmented_objects_dir, f'{object_id}.png') obj_img.save(object_image_path) metadata = { 'object_id': object_id, 'master_id': master_id, 'object_image_path': object_image_path } object_metadata.append(metadata) metadata_file = os.path.join(segmented_objects_dir, f'{master_id}_metadata.json') with open(metadata_file, 'w') as f: json.dump(object_metadata, f, indent=4) return object_metadata # Run inference #print(outputs) # This will print the model's output, including masks, labels, and scores # def extract_objects(image, masks): # """ # Extract objects from the segmented image using masks. # Args: # - image (PIL.Image): The original image. # - masks (Tensor): Masks obtained from the segmentation model. # Returns: # - List of extracted objects as images. # """ # image_np = np.array(image) # extracted_objects = [] # for i, mask in enumerate(masks): # # Convert mask to binary # binary_mask = mask[0].mul(255).byte().cpu().numpy() # # Extract object using the mask # masked_image = cv2.bitwise_and(image_np, image_np, mask=binary_mask) # # Find the bounding box of the object # x, y, w, h = cv2.boundingRect(binary_mask) # cropped_object = masked_image[y:y+h, x:x+w] # # Convert cropped object back to PIL Image # cropped_object_pil = Image.fromarray(cropped_object) # extracted_objects.append(cropped_object_pil) # return extracted_objects # import os # import uuid # from PIL import Image # import json # # Directories to save the input images and segmented objects # input_images_dir = 'data/input_images/' # segmented_objects_dir = 'data/segmented_objects/' # os.makedirs(input_images_dir, exist_ok=True) # os.makedirs(segmented_objects_dir, exist_ok=True) # def save_input_image(image, master_id): # """ # Save the original input image with a unique master ID. # Args: # - image (PIL.Image): The original input image. # - master_id (str): Unique ID for the original image. # Returns: # - str: Path to the saved input image. # """ # input_image_path = os.path.join(input_images_dir, f'{master_id}.png') # image.save(input_image_path) # return input_image_path # def save_objects_and_metadata(extracted_objects, master_id): # """ # Save the extracted objects as images and store their metadata. # Args: # - extracted_objects (List[PIL.Image]): List of extracted objects as images. # - master_id (str): Unique ID for the original image. # Returns: # - List of metadata dictionaries for each object. # """ # object_metadata = [] # for i, obj_img in enumerate(extracted_objects): # # Generate a unique ID for each object # object_id = str(uuid.uuid4()) # # Save the object image # object_image_path = os.path.join(segmented_objects_dir, f'{object_id}.png') # obj_img.save(object_image_path) # # Prepare metadata for the object # metadata = { # 'object_id': object_id, # 'master_id': master_id, # 'object_image_path': object_image_path # } # object_metadata.append(metadata) # # Save metadata to JSON (or you can save to a database) # metadata_file = os.path.join(segmented_objects_dir, f'{master_id}_metadata.json') # with open(metadata_file, 'w') as f: # json.dump(object_metadata, f, indent=4) # return object_metadata # # Example usage # master_id = str(uuid.uuid4()) # Generate a unique master ID for the original image # # Save the input image # input_image_path = save_input_image(image, master_id) # # Save the objects and their metadata # metadata = save_objects_and_metadata(extracted_objects, master_id) # import cv2 # import os # import json # import uuid # import numpy as np # from PIL import Image # # Directories to save the segmented objects and metadata # segmented_objects_dir = 'data/segmented_objects/' # metadata_file = 'data/segmented_objects_metadata.json' # # Ensure directories exist # os.makedirs(segmented_objects_dir, exist_ok=True) # def extract_objects(image_path, masks, master_id): # # Load the original image # image = Image.open(image_path) # image_np = np.array(image) # object_metadata = [] # for i, mask in enumerate(masks): # # Generate a unique ID for each object # object_id = str(uuid.uuid4()) # # Extract object using the mask # masked_image = cv2.bitwise_and(image_np, image_np, mask=mask) # # Find the bounding box of the object # x, y, w, h = cv2.boundingRect(mask) # cropped_object = masked_image[y:y+h, x:x+w] # # Save the object image # object_image_path = os.path.join(segmented_objects_dir, f'{object_id}.png') # cv2.imwrite(object_image_path, cropped_object) # # Save metadata # object_metadata.append({ # 'object_id': object_id, # 'master_id': master_id, # 'object_image_path': object_image_path, # 'bounding_box': (x, y, w, h) # }) # # Save metadata to JSON # with open(metadata_file, 'w') as f: # json.dump(object_metadata, f, indent=4) # return object_metadata # # Example usage: # # Assuming `masks` is a list of binary masks (numpy arrays) from your segmentation model # # and `image_path` is the path to the original image # master_id = str(uuid.uuid4()) # image_path = 'data/input_images/sample_image.png' # masks = [...] # Replace with actual masks # object_metadata = extract_objects(image_path, masks, master_id) # #Extracting and saving segmented objects # # def save_segmented_objects(image_path, outputs, output_dir='data\segmented_objects'): # # image = Image.open(image_path).convert("RGB") # # image_np = np.array(image) # # masks = outputs[0]['masks'] # # scores = outputs[0]['scores'] # # if not os.path.exists(output_dir): # # os.makedirs(output_dir) # # for i in range(len(scores)): # # if scores[i] > 0.5: # Confidence threshold # # mask = masks[i].squeeze().cpu().numpy() # # mask = np.where(mask > 0.5, 1, 0).astype(np.uint8) # Binarize mask # # # Create a new image for the masked object # # masked_image = np.zeros_like(image_np) # # for c in range(3): # Apply the mask to each channel (R, G, B) # # masked_image[:, :, c] = image_np[:, :, c] * mask # # # Save the masked object # # masked_image_pil = Image.fromarray(masked_image) # # masked_image_pil.save(f"{output_dir}object_{i+1}.png") # # # Run the function to save segmented objects # # save_segmented_objects(image_path, outputs)