Spaces:
Sleeping
Sleeping
File size: 1,749 Bytes
01e1043 056020b 01e1043 056020b 01e1043 056020b 01e1043 056020b 01e1043 056020b 01e1043 056020b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 |
import cv2
import datetime
from matplotlib.colors import hsv_to_rgb
import torch
import numpy as np
from super_gradients.training import models
from super_gradients.training.models.detection_models.customizable_detector import CustomizableDetector
from super_gradients.training.pipelines.pipelines import DetectionPipeline
from deep_sort_torch.deep_sort.deep_sort import DeepSort
import os
# make sure to set IOU and confidence in the pipeline constructor
def get_color(number):
""" Converts an integer number to a color """
# change these however you want to
hue = number*30 % 180
saturation = number*103 % 256
value = number*50 % 256
# expects normalized values
hsv_array = [hue/179, saturation/255, value/255]
rgb = hsv_to_rgb(hsv_array)
return [int(c*255) for c in rgb]
def img_predict(media, model, out_path,filename):
save_to = os.path.join(out_path, filename)
images_predictions = model.predict(media,conf=0.70,fuse_model=False)
images_predictions.save(output_folder=out_path, box_thickness=2, show_confidence=True)
return None
def get_prediction(model, image_in, pipeline):
''' Obtains DetectionPrediction object from a single input RGB image
'''
# Preprocess
preprocessed_image, processing_metadata = pipeline.image_processor.preprocess_image(image=image_in.copy())
# Predict
with torch.no_grad():
torch_input = torch.Tensor(preprocessed_image).unsqueeze(0).to('cuda')
model_output = model(torch_input)
prediction = pipeline._decode_model_output(model_output, model_input=torch_input)
# Postprocess
return pipeline.image_processor.postprocess_predictions(predictions=prediction[0], metadata=processing_metadata)
|