Spaces:
Sleeping
Sleeping
| from flask import Flask, request, jsonify | |
| from flask_restx import Api, Resource, fields | |
| from PIL import Image | |
| import numpy as np | |
| import io | |
| import base64 | |
| import pandas as pd | |
| import torch | |
| from transformers import DPTImageProcessor, DPTForDepthEstimation, SegformerImageProcessor, SegformerForSemanticSegmentation | |
| import json | |
| id2label = {} | |
| with open('id2label.json', 'r') as json_file: | |
| id2label = json.load(json_file) | |
| app = Flask(__name__) | |
| api = Api(app, version='1.0', title='Food Segmentation API', description='API for food segmentation and weight estimation') | |
| # Define the image payload model | |
| image_payload = api.model('ImagePayload', { | |
| 'image': fields.String(required=True, description='Base64 encoded image') | |
| }) | |
| # Initialize models | |
| depth_processor = DPTImageProcessor.from_pretrained("Intel/dpt-hybrid-midas") | |
| depth_model = DPTForDepthEstimation.from_pretrained("Intel/dpt-hybrid-midas", low_cpu_mem_usage=True) | |
| seg_processor = SegformerImageProcessor.from_pretrained("prem-timsina/segformer-b0-finetuned-food") | |
| seg_model = SegformerForSemanticSegmentation.from_pretrained("prem-timsina/segformer-b0-finetuned-food") | |
| # Load density database | |
| density_db_path = 'density_table_foodseg103.xlsx' | |
| density_db = pd.read_excel(density_db_path, sheet_name=0, usecols=[1, 2]) | |
| density_db.dropna(inplace=True) | |
| def get_density(food_type): | |
| match = density_db[density_db['Food'].str.contains(food_type, case=False, na=False)] | |
| if not match.empty: | |
| return match.iloc[0]['Density in g/ml'] | |
| return 0 # Default density if not found | |
| def estimate_depth(image): | |
| inputs = depth_processor(images=image, return_tensors="pt") | |
| with torch.no_grad(): | |
| outputs = depth_model(**inputs) | |
| predicted_depth = outputs.predicted_depth | |
| prediction = torch.nn.functional.interpolate( | |
| predicted_depth.unsqueeze(1), | |
| size=image.size[::-1], | |
| mode="bicubic", | |
| align_corners=False, | |
| ) | |
| return prediction.squeeze().cpu().numpy() | |
| def segment_image(image): | |
| inputs = seg_processor(images=image, return_tensors="pt") | |
| with torch.no_grad(): | |
| outputs = seg_model(**inputs) | |
| logits = outputs.logits | |
| return torch.argmax(logits, dim=1).squeeze().cpu().numpy() | |
| def calculate_volume(depth_map, mask): | |
| if depth_map.shape != mask.shape: | |
| depth_map = np.resize(depth_map, mask.shape) | |
| masked_depth = depth_map * mask | |
| volume = round(np.sum(masked_depth) * (0.0001),2) | |
| return volume | |
| class ProcessImage(Resource): | |
| def post(self): | |
| data = request.json | |
| base64_image = data['image'] | |
| # Decode the image | |
| image_data = base64.b64decode(base64_image) | |
| image = Image.open(io.BytesIO(image_data)).convert("RGB") | |
| # Estimate depth and segment image | |
| depth_map = estimate_depth(image) | |
| segmentation_mask = segment_image(image) | |
| # Ensure depth_map and segmentation_mask have the same dimensions | |
| if depth_map.shape != segmentation_mask.shape: | |
| depth_map = np.resize(depth_map, segmentation_mask.shape) | |
| # Calculate volume and weight for each class | |
| unique_classes = np.unique(segmentation_mask) | |
| results = [] | |
| for class_id in unique_classes: | |
| mask = (segmentation_mask == class_id).astype(np.uint8) | |
| volume = calculate_volume(depth_map, mask) | |
| class_name = id2label[str(class_id)] | |
| density = get_density(class_name) | |
| weight = round(volume * density,2) | |
| results.append({ | |
| 'class_id': int(class_id), | |
| 'class_name': class_name, | |
| 'estimated_weight_g': weight | |
| }) | |
| return jsonify(results) | |
| #if __name__ == "__main__": | |
| # app.run(host='0.0.0.0', port=5000) | |