Spaces:
Sleeping
Sleeping
| from detection import ml_detection, ml_utils | |
| import json | |
| from PIL import Image | |
| # Run detection pipeline: load ML model, perform object detection and return json object | |
| def detection_pipeline(model_type, image_bytes): | |
| # Load correct ML model | |
| detr_processor, detr_model = ml_detection.load_model(model_type) | |
| # Perform object detection | |
| results = ml_detection.object_detection(detr_processor, detr_model, image_bytes) | |
| # Convert dictionary of tensors to JSON object | |
| result_json_dict = ml_utils.convert_tensor_dict_to_json(results) | |
| return result_json_dict | |
| def main(): | |
| print('Main function') | |
| model_type = "facebook/detr-resnet-50" | |
| image_path = './samples/boats.jpg' | |
| # Reading image file as image_bytes (similar to API request) | |
| print('Reading image file...') | |
| with open(image_path, 'rb') as image_file: | |
| image_bytes = image_file.read() | |
| result_json = detection_pipeline(model_type, image_bytes) | |
| print("result_json:", result_json) | |
| if __name__ == "__main__": | |
| main() | |