detector.py

from LLM_package import GeminiInference 
import json
class ObjectDetector:
    def __init__(self, model_path):
        self.model  = GeminiInference() 
        self.prompt_objects=None
        self.text=None
        self.prompt= f"""
    Detect all {self.prompt_objects} in the image. The box_2d should be [ymin, xmin, ymax, xmax] normalized to 0-1000.
    Please provide the response as a JSON array of objects, where each object has a 'label' and 'box_2d' field.
    Example:
    [
        {{"label": "face", "box_2d": [100, 200, 300, 400]}},
        {{"label": "license_plate", "box_2d": [500, 600, 700, 800]}}
    ]
    """
    def detect_objects(self, image_path):
        detected_objects_norm_0_1= self.model.parse_response(self.model.get_response(image_path, self.prompt))
        return detected_objects_norm_0_1
    """
    Detects the danger level of the image.
    """
    def detect_danger_level(self, image_path):
        analysis_prompt = f"""
            画像の個人情報漏洩リスクを分析し、厳密にJSON形式で返答してください。なおこの時、資料があれば、資料を参考にしてください：
                {{
                    "risk_level": "high|medium|low",
                    "risk_reason": "リスクの具体的理由",
                    "objects_to_remove": ["消去すべきオブジェクトリスト(英語で、例: 'face', 'license_plate')"]
                }}
                <資料>
                {self.text if self.text else "なし"}
                </資料>
                """
        
        response = self.model.parse(self.model.get_response(image_path, analysis_prompt))
        print(f"Response: {response}")
        return response