add edit

Dockerfile

@@ -35,4 +35,4 @@ USER appuser
 EXPOSE 7860
 
 # Run the server
-CMD ["python", "/app/server.py"]
+CMD ["python", "/app/main.py"]

blender_script.py

@@ -0,0 +1,58 @@
+import bpy
+import os
+import json
+from mathutils import Vector
+
+def clear_scene():
+    bpy.ops.object.select_all(action="SELECT")
+    bpy.ops.object.delete()
+
+def setup_scene(objects, environment):
+    clear_scene()
+    # Load environment (e.g., desert)
+    bpy.ops.wm.open_mainfile(filepath=f"models/{environment}.blend")
+    
+    # Import military assets
+    for obj in objects:
+        bpy.ops.import_scene.obj(filepath=f"models/{obj}.obj")
+    
+    # Setup camera
+    bpy.ops.object.camera_add(location=(0, -10, 5), rotation=(1.0, 0, 0))
+    bpy.context.scene.camera = bpy.context.object
+
+def render_image(output_dir, image_id):
+    bpy.context.scene.render.filepath = os.path.join(output_dir, f"image_{image_id}.png")
+    bpy.ops.render.render(write_still=True)
+    
+    # Generate labels (bounding boxes)
+    labels = []
+    for obj in bpy.data.objects:
+        if obj.type == "MESH":
+            # Project 3D bounds to 2D
+            coords_2d = [bpy.context.scene.camera.matrix_world @ Vector(corner) for corner in obj.bound_box]
+            coords_2d = [bpy_extras.object_utils.world_to_camera_view(bpy.context.scene, bpy.context.scene.camera, coord) for coord in coords_2d]
+            x_coords = [coord.x * bpy.context.scene.render.resolution_x for coord in coords_2d]
+            y_coords = [(1 - coord.y) * bpy.context.scene.render.resolution_y for coord in coords_2d]  # Invert y-axis
+            x_min, x_max = min(x_coords), max(x_coords)
+            y_min, y_max = min(y_coords), max(y_coords)
+            labels.append({
+                "category": obj.name,
+                "bbox": [x_min, y_min, x_max - x_min, y_max - y_min]
+            })
+    return labels
+
+def generate_images(objects, environment, num_images, output_dir):
+    setup_scene(objects, environment)
+    annotations = []
+    for i in range(num_images):
+        # Randomize object positions, camera angles, etc. (simplified here)
+        labels = render_image(output_dir, i)
+        annotations.append({"image_id": i, "file_name": f"image_{i}.png", "labels": labels})
+    
+    with open(os.path.join(output_dir, "annotations.json"), "w") as f:
+        json.dump(annotations, f)
+
+if __name__ == "__main__":
+    import sys
+    objects, env, num, out_dir = sys.argv[1], sys.argv[2], int(sys.argv[3]), sys.argv[4]
+    generate_images(objects.split(","), env, num, out_dir)

main.py

@@ -0,0 +1,91 @@
+from fastapi import FastAPI, HTTPException
+from fastapi.responses import FileResponse
+from pydantic import BaseModel
+import subprocess
+import os
+import shutil
+import tempfile
+import zipfile
+from diffusers import StableDiffusionInstructPix2PixPipeline
+import torch
+from PIL import Image
+import json
+app = FastAPI()
+
+# Load InstructPix2Pix model
+pipe = StableDiffusionInstructPix2PixPipeline.from_pretrained(
+    "timm/instruct-pix2pix",
+    torch_dtype=torch.float16,
+    safety_checker=None,
+).to("cuda")
+
+class DatasetRequest(BaseModel):
+    objects: list[str]
+    environment: str
+    num_images: int
+    augmentation_prompts: list[str]
+
+def augment_image(image_path, prompt):
+    image = Image.open(image_path).convert("RGB")
+    augmented = pipe(prompt=prompt, image=image, num_inference_steps=20, image_guidance_scale=1.5).images[0]
+    return augmented
+
+@app.post("/generate_dataset")
+async def generate_dataset(request: DatasetRequest):
+    try:
+        with tempfile.TemporaryDirectory() as tmpdirname:
+            # Step 1: Generate base images with Blender
+            base_dir = os.path.join(tmpdirname, "base")
+            os.makedirs(base_dir)
+            subprocess.run([
+                "blender", "--background", "--python", "blender_script.py", "--",
+                ",".join(request.objects), request.environment, str(request.num_images), base_dir
+            ], check=True)
+            
+            # Load base annotations
+            with open(os.path.join(base_dir, "annotations.json"), "r") as f:
+                base_annotations = json.load(f)
+            
+            # Step 2: Augment images
+            output_dir = os.path.join(tmpdirname, "output/images")
+            os.makedirs(output_dir)
+            annotations = []
+            image_id = 0
+            for base_anno in base_annotations:
+                base_image_path = os.path.join(base_dir, base_anno["file_name"])
+                for prompt in request.augmentation_prompts:
+                    augmented = augment_image(base_image_path, prompt)
+                    new_filename = f"image_{image_id}.png"
+                    augmented.save(os.path.join(output_dir, new_filename))
+                    annotations.append({
+                        "image_id": image_id,
+                        "file_name": new_filename,
+                        "labels": base_anno["labels"]
+                    })
+                    image_id += 1
+            
+            # Save annotations
+            anno_file = os.path.join(tmpdirname, "output/annotations.json")
+            with open(anno_file, "w") as f:
+                json.dump(annotations, f)
+            
+            # Step 3: Create zip file
+            zip_path = os.path.join(tmpdirname, "dataset.zip")
+            with zipfile.ZipFile(zip_path, "w") as zipf:
+                for root, _, files in os.walk(output_dir):
+                    for file in files:
+                        zipf.write(os.path.join(root, file), os.path.join("images", file))
+                zipf.write(anno_file, "annotations.json")
+            
+            return FileResponse(zip_path, media_type="application/zip", filename="dataset.zip")
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+@app.get("/health")
+async def health_check():
+    return {"status": "healthy"}
+
+
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=7860)

requirements.txt

@@ -4,4 +4,7 @@ torch
 diffusers 
 huggingface_hub 
 safetensors
-transformers
+transformers
+pillow
+numpy
+blenderproc