hlopez/Waste-Detector

1from typing import Tuple2import matplotlib.pyplot as plt3import numpy as np4import cv25import torch6 7from icevision.models.checkpoint import model_from_checkpoint8from classifier import CustomViT9 10def plot_img_no_mask(image : np.ndarray, boxes : torch.Tensor, labels):11    colors = {12        0: (255,255,0),13        1: (255, 0, 0),14        2: (0, 0, 255),15        3: (0,128,0),16        4: (255,165,0),17        5: (230,230,250),18        6: (192,192,192)19    }20 21    texts = {22        0: 'plastic',23        1: 'dangerous',24        2: 'carton',25        3: 'glass',26        4: 'organic',27        5: 'rest',28        6: 'other'29    }30 31    # Show image32    boxes = boxes.cpu().detach().numpy().astype(np.int32)33    fig, ax = plt.subplots(1, 1, figsize=(12, 6))34 35    for i, box in enumerate(boxes):36        color = colors[labels[i]]37 38        [x1, y1, x2, y2] = np.array(box).astype(int)39        # Si no se hace la copia da error en cv2.rectangle40        image = np.array(image).copy()41 42        pt1 = (x1, y1)43        pt2 = (x2, y2)44        cv2.rectangle(image, pt1, pt2, color, thickness=5)45        cv2.putText(image, texts[labels[i]], (x1, y1-10),46                    cv2.FONT_HERSHEY_SIMPLEX, 4, thickness=5, color=color)47 48    plt.axis('off')49    ax.imshow(image)50 51    fig.savefig("img.png", bbox_inches='tight')52 53def get_models(54    detection_ckpt : str,55    classifier_ckpt : str56) -> Tuple[torch.nn.Module, torch.nn.Module]:57    """58    Get the detection and classifier models59 60    Args:61        detection_ckpt (str): Detection model checkpoint62        classifier_ckpt (str): Classifier model checkpoint63 64    Returns:65        tuple: Tuple containing:66            - (torch.nn.Module): Detection model67            - (torch.nn.Module): Classifier model68    """69    print('Loading the detection model')70    checkpoint_and_model = model_from_checkpoint(71                                detection_ckpt,72                                model_name='ross.efficientdet',73                                backbone_name='d0',74                                img_size=512,75                                classes=['Waste'],76                                revise_keys=[(r'^model\.', '')],77                                map_location='cpu')78 79    det_model = checkpoint_and_model['model']80    det_model.eval()81 82    print('Loading the classifier model')83    classifier = CustomViT(target_size=7, pretrained=False)84    classifier.load_state_dict(torch.load(classifier_ckpt, map_location='cpu'))85    classifier.eval()86 87    return det_model, classifier