from __future__ import annotations import os import io from uuid import uuid4 from typing import List, Tuple from django.core.management.base import BaseCommand, CommandError from django.core.files import File from django.contrib.auth import get_user_model from PIL import Image from anatomy.models import AnatomyQuestion, Exam, ExamQuestionDetail, QuestionType, Structure from generic.models import Modality as GenericModality # type: ignore class Command(BaseCommand): help = ( "Create anatomy questions and assign them to an exam." "Requires --exam-id and --image-path (path readable inside the container)." ) def add_arguments(self, parser): parser.add_argument("--exam-id", type=int, required=False, help="ID of the anatomy.Exam to assign questions to (if omitted a new exam will be created)") parser.add_argument("--exam-name", type=str, default=None, help="Name to use when creating a new exam if --exam-id omitted") parser.add_argument("--count", type=int, default=1, help="Number of questions to create") parser.add_argument("--image-path", type=str, required=False, help="Path to image file to use for each created question") parser.add_argument("--generate-image", action="store_true", help="Generate a placeholder image automatically if --image-path is not provided or not found") parser.add_argument("--image-size", type=str, default="800x600", help="Generated image size WIDTHxHEIGHT, e.g. 800x600") parser.add_argument("--image-color", type=str, default="#777777", help="Background color for generated image (hex) e.g. #ffffff") parser.add_argument("--description", type=str, default="Generated question {i}", help="Description template; use {i} for index") parser.add_argument("--question-type-id", type=int, default=None, help="Optional QuestionType id to assign") parser.add_argument("--modality-id", type=int, default=None, help="Optional Modality id to assign (generic.Modality id)") parser.add_argument("--structure-id", type=int, default=None, help="Optional Structure id to assign") parser.add_argument("--author-id", type=int, action="append", help="Author user id; repeatable") parser.add_argument("--start-sort", type=int, default=1, help="Starting sort_order for the first question") def handle(self, *args, **options): exam_id = options.get("exam_id") count = options["count"] image_path = options.get("image_path") generate_image = options.get("generate_image") image_size_str = options.get("image_size") or "800x600" image_color = options.get("image_color") or "#777777" description_template = options["description"] qtype_id = options.get("question_type_id") modality_id = options.get("modality_id") structure_id = options.get("structure_id") author_ids: List[int] = options.get("author_id") or [] sort = options.get("start_sort") or 1 exam = None if exam_id: try: exam = Exam.objects.get(pk=exam_id) except Exam.DoesNotExist: raise CommandError(f"Exam with id {exam_id} does not exist") else: # create new exam exam_name = options.get("exam_name") or f"Auto Anatomy Exam {uuid4().hex[:8]}" exam = Exam(name=exam_name, active=True, open_access=True) exam.save() self.stdout.write(self.style.SUCCESS(f"Created exam id={exam.id} name='{exam.name}'")) # validate image path if provided if image_path and not os.path.exists(image_path): if not generate_image: raise CommandError(f"Image file not found at path: {image_path}") else: image_path = None # Resolve optional foreign keys qtype = None if qtype_id: qtype = QuestionType.objects.filter(pk=qtype_id).first() if not qtype: raise CommandError(f"QuestionType id {qtype_id} not found") modality = None if modality_id: # The project stores Modality in generic.models.Modality try: from generic.models import Modality as GModality modality = GModality.objects.filter(pk=modality_id).first() if not modality: raise CommandError(f"Modality id {modality_id} not found in generic.Modality") except Exception as e: raise CommandError(str(e)) structure = None if structure_id: from atlas.models import Structure as AtlasStructure structure = AtlasStructure.objects.filter(pk=structure_id).first() if not structure: raise CommandError(f"Structure id {structure_id} not found") User = get_user_model() authors = [] for uid in author_ids: u = User.objects.filter(pk=uid).first() if not u: raise CommandError(f"Author user id {uid} not found") authors.append(u) # parse image size try: w, h = image_size_str.lower().split("x") image_size: Tuple[int, int] = (int(w), int(h)) except Exception: raise CommandError("Invalid --image-size; expected WIDTHxHEIGHT, e.g. 800x600") created = [] for i in range(1, count + 1): desc = description_template.format(i=i) # determine image source: either provided path or generated image unique_name = f"{uuid4().hex}_image.jpg" if image_path: # use provided file with open(image_path, "rb") as f: django_file = File(f, name=unique_name) img_file_for_save = django_file # create and save question q = AnatomyQuestion( question_type=qtype, description=desc, answer_help="", ) if modality: q.modality = modality if structure: q.structure = structure q.image.save(unique_name, img_file_for_save, save=False) q.save() else: # generate an image in memory buf = io.BytesIO() img = Image.new("RGB", image_size, image_color) img.save(buf, format="JPEG") buf.seek(0) django_file = File(buf, name=unique_name) q = AnatomyQuestion( question_type=qtype, description=desc, answer_help="", ) if modality: q.modality = modality if structure: q.structure = structure q.image.save(unique_name, django_file, save=False) q.save() # authors if authors: q.author.set(authors) # add to exam via the many-to-many through relationship try: exam.exam_questions.add(q, through_defaults={"sort_order": sort}) except TypeError: ExamQuestionDetail.objects.create(exam=exam, anatomyquestion=q, sort_order=sort) self.stdout.write(self.style.SUCCESS(f"Created question id={q.id} and added to exam id={exam.id} with sort_order={sort}")) created.append(q) sort += 1 self.stdout.write(self.style.SUCCESS(f"Finished: created {len(created)} question(s)"))