<reponame>meryusha/seeds_faster<filename>maskrcnn_benchmark/data/datasets/evaluation/seed/seed_predict.py
# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
import cv2
import torch
from torchvision import transforms as T

from maskrcnn_benchmark.modeling.detector import build_detection_model
from maskrcnn_benchmark.utils.checkpoint import DetectronCheckpointer
from maskrcnn_benchmark.structures.image_list import to_image_list
from maskrcnn_benchmark import layers as L
from maskrcnn_benchmark.utils import cv2_util
import numpy as np

class SeedPredict(object):
    CLASSES_STRAT1 = (
        "__background__ ",
        "germinated",
        "non-germinated"
        # "seed",
        # "radical"
    )
    CLASSES_STRAT2 = (
        "__background__ ",
        # "germinated",
        # "non-germinated"
        "seed",
        "radical"
    )
    def __init__(
        self,
        cfg,
        confidence_threshold=0.7,
        min_image_size=800,
    ):
        self.cfg = cfg.clone()
        self.model = build_detection_model(cfg)
        self.model.eval()
        self.device = torch.device(cfg.MODEL.DEVICE)
        self.model.to(self.device)
        self.min_image_size = min_image_size

        save_dir = cfg.OUTPUT_DIR
        checkpointer = DetectronCheckpointer(cfg, self.model, save_dir=save_dir)
        _ = checkpointer.load(cfg.MODEL.WEIGHT)

        self.transforms = self.build_transform()


        self.cpu_device = torch.device("cpu")
        self.confidence_threshold = confidence_threshold

    def build_transform(self):
        """
        Creates a basic transformation that was used to train the models
        """
        cfg = self.cfg

        # we are loading images with OpenCV, so we don't need to convert them
        # to BGR, they are already! So all we need to do is to normalize
        # by 255 if we want to convert to BGR255 format, or flip the channels
        # if we want it to be in RGB in [0-1] range.
        if cfg.INPUT.TO_BGR255:
            to_bgr_transform = T.Lambda(lambda x: x * 255)
        else:
            to_bgr_transform = T.Lambda(lambda x: x[[2, 1, 0]])

        normalize_transform = T.Normalize(
            mean=cfg.INPUT.PIXEL_MEAN, std=cfg.INPUT.PIXEL_STD
        )

        transform = T.Compose(
            [
                T.ToPILImage(),
                T.Resize(self.min_image_size),
                T.ToTensor(),
                to_bgr_transform,
                normalize_transform,
            ]
        )
        return transform

    def run_on_opencv_image(self, image_bytes):
        """
        Arguments:
            image_path (String): path to an image
            image_bytes

        Returns:
            prediction (BoxList): the detected objects. Additional information
                of the detection properties can be found in the fields of
                the BoxList via `prediction.fields()`
        """
        
        # image = cv2.imread(image_path)
        image = cv2.imdecode(np.frombuffer(image_bytes, np.uint8), -1)
        cv2.imwrite('image.jpg', image)
        predictions = self.compute_prediction(image)

        # image_copy = image.copy()
        # boxes = predictions.bbox.numpy()
        # print( predictions.get_field("scores"))
        #Checking
        # for box in boxes:
        #     minX = box[0]
        #     minY = box[1]
        #     maxX = box[2]
        #     maxY = box[3]
            # cv2.rectangle(image_copy, ( minX, minY), (maxX, maxY), 0, 4)
        # cv2.imwrite('image.jpg', image_copy)
        # top_predictions = self.select_top_predictions(predictions)
        # predictions = []
        # result = image.copy()


        return predictions

    def compute_prediction(self, original_image):
        """
        Arguments:
            original_image (np.ndarray): an image as returned by OpenCV

        Returns:
            prediction (BoxList): the detected objects. Additional information
                of the detection properties can be found in the fields of
                the BoxList via `prediction.fields()`
        """
        # apply pre-processing to image
        image = self.transforms(original_image)
        # convert to an ImageList, padded so that it is divisible by
        # cfg.DATALOADER.SIZE_DIVISIBILITY
        image_list = to_image_list(image, self.cfg.DATALOADER.SIZE_DIVISIBILITY)
        image_list = image_list.to(self.device)
        # compute predictions
        print("ALL fine here 1")
        with torch.no_grad():
            predictions = self.model(image_list)
        print("ALL fine here 2")
        predictions = [o.to(self.cpu_device) for o in predictions]

        # always single image is passed at a time
        prediction = predictions[0]

        # reshape prediction (a BoxList) into the original image size
        height, width = original_image.shape[:-1]
        prediction = prediction.resize((width, height))
    
        return prediction

    def select_top_predictions(self, predictions):
        """
        Select only predictions which have a `score` > self.confidence_threshold,
        and returns the predictions in descending order of score

        Arguments:
            predictions (BoxList): the result of the computation by the model.
                It should contain the field `scores`.

        Returns:
            prediction (BoxList): the detected objects. Additional information
                of the detection properties can be found in the fields of
                the BoxList via `prediction.fields()`
        """
        scores = predictions.get_field("scores")
        keep = torch.nonzero(scores > self.confidence_threshold).squeeze(1)
        predictions = predictions[keep]
        scores = predictions.get_field("scores")
        _, idx = scores.sort(0, descending=True)
        return predictions[idx]


    def overlay_boxes(self, image, predictions):
        """
        Adds the predicted boxes on top of the image

        Arguments:
            image (np.ndarray): an image as returned by OpenCV
            predictions (BoxList): the result of the computation by the model.
                It should contain the field `labels`.
        """
        labels = predictions.get_field("labels")
        boxes = predictions.bbox

        colors = self.compute_colors_for_labels(labels).tolist()

        for box, color in zip(boxes, colors):
            box = box.to(torch.int64)
            top_left, bottom_right = box[:2].tolist(), box[2:].tolist()
            image = cv2.rectangle(
                image, tuple(top_left), tuple(bottom_right), tuple(color), 1
            )

        return image
        
    def map_class_id_to_class_name(self, class_id):  
        return self.CLASSES_STRAT2[class_id]
       

# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
import cv2
import torch
from torchvision import transforms as T

from maskrcnn_benchmark.modeling.detector import build_detection_model
from maskrcnn_benchmark.utils.checkpoint import DetectronCheckpointer
from maskrcnn_benchmark.structures.image_list import to_image_list
from maskrcnn_benchmark import layers as L
from maskrcnn_benchmark.utils import cv2_util
import numpy as np

class SeedPredict(object):
    CLASSES_STRAT1 = (
        "__background__ ",
        "germinated",
        "non-germinated"
        # "seed",
        # "radical"
    )
    CLASSES_STRAT2 = (
        "__background__ ",
        # "germinated",
        # "non-germinated"
        "seed",
        "radical"
    )
    def __init__(
        self,
        cfg,
        confidence_threshold=0.7,
        min_image_size=800,
    ):
        self.cfg = cfg.clone()
        self.model = build_detection_model(cfg)
        self.model.eval()
        self.device = torch.device(cfg.MODEL.DEVICE)
        self.model.to(self.device)
        self.min_image_size = min_image_size

        save_dir = cfg.OUTPUT_DIR
        checkpointer = DetectronCheckpointer(cfg, self.model, save_dir=save_dir)
        _ = checkpointer.load(cfg.MODEL.WEIGHT)

        self.transforms = self.build_transform()


        self.cpu_device = torch.device("cpu")
        self.confidence_threshold = confidence_threshold

    def build_transform(self):
        """
        Creates a basic transformation that was used to train the models
        """
        cfg = self.cfg

        # we are loading images with OpenCV, so we don't need to convert them
        # to BGR, they are already! So all we need to do is to normalize
        # by 255 if we want to convert to BGR255 format, or flip the channels
        # if we want it to be in RGB in [0-1] range.
        if cfg.INPUT.TO_BGR255:
            to_bgr_transform = T.Lambda(lambda x: x * 255)
        else:
            to_bgr_transform = T.Lambda(lambda x: x[[2, 1, 0]])

        normalize_transform = T.Normalize(
            mean=cfg.INPUT.PIXEL_MEAN, std=cfg.INPUT.PIXEL_STD
        )

        transform = T.Compose(
            [
                T.ToPILImage(),
                T.Resize(self.min_image_size),
                T.ToTensor(),
                to_bgr_transform,
                normalize_transform,
            ]
        )
        return transform

    def run_on_opencv_image(self, image_bytes):
        """
        Arguments:
            image_path (String): path to an image
            image_bytes

        Returns:
            prediction (BoxList): the detected objects. Additional information
                of the detection properties can be found in the fields of
                the BoxList via `prediction.fields()`
        """
        
        # image = cv2.imread(image_path)
        image = cv2.imdecode(np.frombuffer(image_bytes, np.uint8), -1)
        cv2.imwrite('image.jpg', image)
        predictions = self.compute_prediction(image)

        # image_copy = image.copy()
        # boxes = predictions.bbox.numpy()
        # print( predictions.get_field("scores"))
        #Checking
        # for box in boxes:
        #     minX = box[0]
        #     minY = box[1]
        #     maxX = box[2]
        #     maxY = box[3]
            # cv2.rectangle(image_copy, ( minX, minY), (maxX, maxY), 0, 4)
        # cv2.imwrite('image.jpg', image_copy)
        # top_predictions = self.select_top_predictions(predictions)
        # predictions = []
        # result = image.copy()


        return predictions

    def compute_prediction(self, original_image):
        """
        Arguments:
            original_image (np.ndarray): an image as returned by OpenCV

        Returns:
            prediction (BoxList): the detected objects. Additional information
                of the detection properties can be found in the fields of
                the BoxList via `prediction.fields()`
        """
        # apply pre-processing to image
        image = self.transforms(original_image)
        # convert to an ImageList, padded so that it is divisible by
        # cfg.DATALOADER.SIZE_DIVISIBILITY
        image_list = to_image_list(image, self.cfg.DATALOADER.SIZE_DIVISIBILITY)
        image_list = image_list.to(self.device)
        # compute predictions
        print("ALL fine here 1")
        with torch.no_grad():
            predictions = self.model(image_list)
        print("ALL fine here 2")
        predictions = [o.to(self.cpu_device) for o in predictions]

        # always single image is passed at a time
        prediction = predictions[0]

        # reshape prediction (a BoxList) into the original image size
        height, width = original_image.shape[:-1]
        prediction = prediction.resize((width, height))
    
        return prediction

    def select_top_predictions(self, predictions):
        """
        Select only predictions which have a `score` > self.confidence_threshold,
        and returns the predictions in descending order of score

        Arguments:
            predictions (BoxList): the result of the computation by the model.
                It should contain the field `scores`.

        Returns:
            prediction (BoxList): the detected objects. Additional information
                of the detection properties can be found in the fields of
                the BoxList via `prediction.fields()`
        """
        scores = predictions.get_field("scores")
        keep = torch.nonzero(scores > self.confidence_threshold).squeeze(1)
        predictions = predictions[keep]
        scores = predictions.get_field("scores")
        _, idx = scores.sort(0, descending=True)
        return predictions[idx]


    def overlay_boxes(self, image, predictions):
        """
        Adds the predicted boxes on top of the image

        Arguments:
            image (np.ndarray): an image as returned by OpenCV
            predictions (BoxList): the result of the computation by the model.
                It should contain the field `labels`.
        """
        labels = predictions.get_field("labels")
        boxes = predictions.bbox

        colors = self.compute_colors_for_labels(labels).tolist()

        for box, color in zip(boxes, colors):
            box = box.to(torch.int64)
            top_left, bottom_right = box[:2].tolist(), box[2:].tolist()
            image = cv2.rectangle(
                image, tuple(top_left), tuple(bottom_right), tuple(color), 1
            )

        return image
        
    def map_class_id_to_class_name(self, class_id):  
        return self.CLASSES_STRAT2[class_id]
# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. # "seed", # "radical" # "germinated", # "non-germinated" Creates a basic transformation that was used to train the models # we are loading images with OpenCV, so we don't need to convert them # to BGR, they are already! So all we need to do is to normalize # by 255 if we want to convert to BGR255 format, or flip the channels # if we want it to be in RGB in [0-1] range. Arguments:             image_path (String): path to an image             image_bytes          Returns:             prediction (BoxList): the detected objects. Additional information                 of the detection properties can be found in the fields of                 the BoxList via `prediction.fields()` # image = cv2.imread(image_path) # image_copy = image.copy() # boxes = predictions.bbox.numpy() # print( predictions.get_field("scores")) #Checking # for box in boxes: #     minX = box[0] #     minY = box[1] #     maxX = box[2] #     maxY = box[3] # cv2.rectangle(image_copy, ( minX, minY), (maxX, maxY), 0, 4) # cv2.imwrite('image.jpg', image_copy) # top_predictions = self.select_top_predictions(predictions) # predictions = [] # result = image.copy() Arguments:             original_image (np.ndarray): an image as returned by OpenCV          Returns:             prediction (BoxList): the detected objects. Additional information                 of the detection properties can be found in the fields of                 the BoxList via `prediction.fields()` # apply pre-processing to image # convert to an ImageList, padded so that it is divisible by # cfg.DATALOADER.SIZE_DIVISIBILITY # compute predictions # always single image is passed at a time # reshape prediction (a BoxList) into the original image size Select only predictions which have a `score` > self.confidence_threshold,         and returns the predictions in descending order of score          Arguments:             predictions (BoxList): the result of the computation by the model.                 It should contain the field `scores`.          Returns:             prediction (BoxList): the detected objects. Additional information                 of the detection properties can be found in the fields of                 the BoxList via `prediction.fields()` Adds the predicted boxes on top of the image          Arguments:             image (np.ndarray): an image as returned by OpenCV             predictions (BoxList): the result of the computation by the model.                 It should contain the field `labels`.