Merge main into interpretability-clean

shivv23 · shivv23 · commit 9bf2dc4602cd · 2026-04-26T23:49:06.000+05:30
diff --git a/tensormap-backend/app/routers/data_process.py b/tensormap-backend/app/routers/data_process.py
@@ -11,6 +11,7 @@
 )
 from app.services.data_process import (
     add_target_service,
+    augment_image_service,
     delete_one_target_by_id_service,
     get_all_targets_service,
     get_column_stats_service,
@@ -95,3 +96,28 @@ def preprocess(file_id: uuid_pkg.UUID, request: PreprocessRequest, db: Session =
     logger.debug("Preprocessing file_id=%s with %d transformations", file_id, len(request.transformations))
     body, status_code = preprocess_data(db, file_id=file_id, transformations=request.transformations)
     return JSONResponse(status_code=status_code, content=body)
+
+
+@router.post("/data/augment/image/{file_id}")
+def augment_image(
+    file_id: uuid_pkg.UUID,
+    technique: str = Query(
+        "flip_horizontal",
+        pattern="^(flip_horizontal|flip_vertical|rotate_90|brightness|zoom|gaussian_noise|random_crop)$",
+    ),
+    db: Session = Depends(get_db),
+):
+    """Apply image augmentation techniques to generate synthetic variants.
+
+    Supported techniques:
+    - flip_horizontal: Mirror along vertical axis
+    - flip_vertical: Mirror along horizontal axis
+    - rotate_90: Rotate by 90 degrees
+    - brightness: Increase brightness by 20%
+    - zoom: Zoom to 90% then resize
+    - gaussian_noise: Add Gaussian noise
+    - random_crop: Crop to 85% then resize
+    """
+    logger.debug("Applying %s augmentation to file_id=%s", technique, file_id)
+    body, status_code = augment_image_service(db, file_id=file_id, technique=technique)
+    return JSONResponse(status_code=status_code, content=body)
diff --git a/tensormap-backend/app/routers/deep_learning.py b/tensormap-backend/app/routers/deep_learning.py
@@ -10,6 +10,7 @@
 from app.schemas.deep_learning import ModelNameRequest, ModelSaveRequest, ModelValidateRequest, TrainingConfigRequest
 from app.services.deep_learning import (
     delete_model_service,
+    export_model_service,
     get_available_model_list,
     get_code_service,
     get_model_graph_service,
@@ -125,3 +126,16 @@ def interpret_model(
     logger.debug("Interpreting model %s", model_name)
     body, status_code = interpret_model_service(db, model_name=model_name, file_id=file_id, project_id=project_id)
     return JSONResponse(status_code=status_code, content=body)
+
+
+@router.get("/model/export/{model_name}")
+def export_model(
+    model_name: str,
+    format: str = Query("savedmodel", pattern="^(savedmodel|tflite|onnx)$"),
+    project_id: uuid_pkg.UUID | None = Query(None),
+    db: Session = Depends(get_db),
+):
+    """Export a trained model in the specified format."""
+    logger.debug("Exporting model %s as %s", model_name, format)
+    body, status_code = export_model_service(db, model_name=model_name, export_format=format, project_id=project_id)
+    return JSONResponse(status_code=status_code, content=body)
diff --git a/tensormap-backend/app/services/data_process.py b/tensormap-backend/app/services/data_process.py
@@ -1,3 +1,4 @@
+import os
 import uuid as uuid_pkg
 from collections.abc import Callable
 from typing import Any
@@ -370,3 +371,86 @@ def preprocess_data(db: Session, file_id: uuid_pkg.UUID, transformations: list)
     except Exception as e:
         logger.exception("Error preprocessing data: %s", str(e))
         return _resp(500, False, f"Error preprocessing data: {e}")
+
+
+def augment_image_service(
+    db: Session,
+    file_id: uuid_pkg.UUID,
+    technique: str = "flip_horizontal",
+) -> tuple:
+    """Apply image augmentation techniques to generate synthetic variants.
+
+    Supported techniques:
+    - flip_horizontal: Mirror image along vertical axis
+    - flip_vertical: Mirror image along horizontal axis
+    - rotate_90: Rotate image by 90 degrees
+    - brightness: Adjust brightness by 20%
+    - zoom: Zoom to 90% then resize
+    - gaussian_noise: Add Gaussian noise
+    - random_crop: Crop to 85% then resize
+    """
+    file_record = db.get(DataFile, file_id)
+    if not file_record:
+        return _resp(404, False, "Image file not found")
+
+    file_path = file_record.file_path
+    if not file_path or not os.path.exists(file_path):
+        return _resp(404, False, "Image file not found on disk")
+
+    settings = get_settings()
+    output_dir = os.path.join(settings.UPLOAD_DIRECTORY, f"augmented_{file_id}")
+    os.makedirs(output_dir, exist_ok=True)
+
+    supported_formats = {".jpg", ".jpeg", ".png", ".bmp", ".gif", ".webp"}
+    _, ext = os.path.splitext(file_path)
+    if ext.lower() not in supported_formats:
+        return _resp(400, False, f"Unsupported image format: {ext}")
+
+    try:
+        from PIL import Image, ImageEnhance
+
+        original = Image.open(file_path)
+
+        if technique == "flip_horizontal":
+            augmented = original.transpose(Image.FLIP_LEFT_RIGHT)
+        elif technique == "flip_vertical":
+            augmented = original.transpose(Image.FLIP_TOP_BOTTOM)
+        elif technique == "rotate_90":
+            augmented = original.rotate(90, expand=True)
+        elif technique == "brightness":
+            enhancer = ImageEnhance.Brightness(original)
+            augmented = enhancer.enhance(1.2)
+        elif technique == "zoom":
+            width, height = original.size
+            new_width = int(width * 0.9)
+            new_height = int(height * 0.9)
+            left = (width - new_width) // 2
+            top = (height - new_height) // 2
+            cropped = original.crop((left, top, left + new_width, top + new_height))
+            augmented = cropped.resize((width, height), Image.LANCZOS)
+        elif technique == "gaussian_noise":
+            np_img = np.array(original.convert("RGB")).astype(np.float32) / 255.0
+            noise = np.random.normal(0, 0.05, np_img.shape)
+            np_img = np.clip(np_img + noise, 0, 1)
+            augmented = Image.fromarray((np_img * 255).astype(np.uint8))
+        elif technique == "random_crop":
+            width, height = original.size
+            crop_size = int(min(width, height) * 0.85)
+            left = np.random.randint(0, width - crop_size + 1)
+            top = np.random.randint(0, height - crop_size + 1)
+            cropped = original.crop((left, top, left + crop_size, top + crop_size))
+            augmented = cropped.resize((width, height), Image.LANCZOS)
+        else:
+            return _resp(400, False, f"Unknown technique: {technique}")
+
+        output_path = os.path.join(output_dir, f"augmented_0{ext}")
+        augmented.save(output_path)
+
+        logger.info("Applied %s augmentation to file %s", technique, file_id)
+        return _resp(200, True, f"Generated augmented image using {technique}", {"output_path": output_path})
+
+    except ImportError:
+        return _resp(500, False, "Pillow not installed")
+    except Exception as e:
+        logger.exception("Error augmenting image: %s", str(e))
+        return _resp(500, False, f"Error augmenting image: {e}")
diff --git a/tensormap-backend/app/services/deep_learning.py b/tensormap-backend/app/services/deep_learning.py
@@ -601,3 +601,63 @@ def interpret_model_service(
     except Exception as e:
         logger.exception("Interpretability failed: %s", str(e))
         return {"success": False, "message": f"Interpretability failed: {e}", "data": None}, 500
+
+
+def export_model_service(
+    db: Session, model_name: str, export_format: str = "savedmodel", project_id: uuid_pkg.UUID | None = None
+) -> tuple:
+    """Export a trained model to various formats.
+
+    Supports: savedmodel (TensorFlow), tflite (TensorFlow Lite), onnx (ONNX).
+    """
+    from app.models import ModelBasic
+
+    stmt = select(ModelBasic).where(ModelBasic.model_name == model_name)
+    if project_id is not None:
+        stmt = stmt.where(ModelBasic.project_id == project_id)
+    model = db.exec(stmt).first()
+
+    if not model:
+        return {"success": False, "message": f"Model '{model_name}' not found", "data": None}, 404
+
+    model_path = os.path.join(MODEL_GENERATION_LOCATION, model_name + MODEL_GENERATION_TYPE)
+    if not os.path.exists(model_path):
+        return {"success": False, "message": "Model file not found", "data": None}, 404
+
+    try:
+        loaded_model = tf.keras.models.load_model(model_path)
+    except Exception as e:
+        logger.error("Failed to load model: %s", e)
+        return {"success": False, "message": f"Could not load model: {e}", "data": None}, 400
+
+    export_dir = os.path.join(MODEL_GENERATION_LOCATION, f"{model_name}_export", export_format)
+    os.makedirs(export_dir, exist_ok=True)
+
+    try:
+        if export_format == "savedmodel":
+            saved_path = os.path.join(export_dir, model_name)
+            loaded_model.save(saved_path)
+            return {"success": True, "message": f"Exported to {saved_path}", "data": {"path": saved_path}}, 200
+
+        if export_format == "tflite":
+            converter = tf.lite.TFLiteConverter.from_keras_model(loaded_model)
+            converter.optimizations = [tf.lite.Optimize.DEFAULT]
+            tflite_model = converter.convert()
+            tflite_path = os.path.join(export_dir, f"{model_name}.tflite")
+            with open(tflite_path, "wb") as f:
+                f.write(tflite_model)
+            return {"success": True, "message": f"Exported to {tflite_path}", "data": {"path": tflite_path}}, 200
+
+        if export_format == "onnx":
+            try:
+                import tf2onnx
+            except ImportError:
+                return {"success": False, "message": "ONNX not available", "data": None}, 501
+            onnx_path = os.path.join(export_dir, f"{model_name}.onnx")
+            tf2onnx.convert.from_keras(loaded_model, output_path=onnx_path)
+            return {"success": True, "message": f"Exported to {onnx_path}", "data": {"path": onnx_path}}, 200
+
+        return {"success": False, "message": f"Unsupported: {export_format}"}, 400
+    except Exception as e:
+        logger.error("Export failed: %s", e)
+        return {"success": False, "message": f"Export failed: {e}"}, 500
diff --git a/tensormap-backend/tests/test_export_service.py b/tensormap-backend/tests/test_export_service.py
@@ -0,0 +1,19 @@
+"""Unit tests for model export service."""
+
+import sys
+from unittest.mock import MagicMock
+
+sys.modules.setdefault("tensorflow", MagicMock())
+sys.modules.setdefault("flatten_json", MagicMock())
+
+
+class TestExportModelService:
+    def test_import(self):
+        from app.services.deep_learning import export_model_service
+
+        assert callable(export_model_service)
+
+    def test_router_import(self):
+        from app.routers.deep_learning import export_model
+
+        assert callable(export_model)
diff --git a/tensormap-backend/tests/test_image_augment.py b/tensormap-backend/tests/test_image_augment.py
@@ -0,0 +1,20 @@
+"""Unit tests for image augmentation service."""
+
+import sys
+from unittest.mock import MagicMock
+
+sys.modules.setdefault("tensorflow", MagicMock())
+sys.modules.setdefault("flatten_json", MagicMock())
+sys.modules.setdefault("pandas", MagicMock())
+
+
+class TestImageAugmentation:
+    def test_import(self):
+        from app.services.data_process import augment_image_service
+
+        assert callable(augment_image_service)
+
+    def test_router_import(self):
+        from app.routers.data_process import augment_image
+
+        assert callable(augment_image)
