mirror of
https://github.com/leaders-of-digital-9-task/backend.git
synced 2024-11-25 10:43:46 +03:00
122 lines
4.1 KiB
Python
122 lines
4.1 KiB
Python
import glob
|
|
import os
|
|
import shutil
|
|
import zipfile
|
|
from pathlib import Path
|
|
|
|
import magic
|
|
import numpy as np
|
|
import pydicom
|
|
from dicom import tasks
|
|
from dicom.models import Coordinate, Dicom, Project
|
|
from django.core.files import File
|
|
from django.core.files.uploadedfile import InMemoryUploadedFile, TemporaryUploadedFile
|
|
from skimage import measure
|
|
from stl import mesh
|
|
from utils.generators import generate_charset
|
|
|
|
|
|
def process_files(
|
|
files: list[TemporaryUploadedFile | InMemoryUploadedFile], user, slug=None
|
|
):
|
|
if slug:
|
|
project = Project.objects.get(slug=slug)
|
|
else:
|
|
project = Project.objects.create(user=user)
|
|
for file in files:
|
|
content_type = magic.from_file(file.temporary_file_path())
|
|
if content_type == "DICOM medical imaging data":
|
|
Dicom.objects.create(file=file, project=project, user=user)
|
|
elif "Zip" in content_type:
|
|
dit_path = f"/tmp/{generate_charset(10)}"
|
|
os.mkdir(dit_path)
|
|
with zipfile.ZipFile(file.temporary_file_path(), "r") as zip_ref:
|
|
zip_ref.extractall(dit_path)
|
|
files = glob.glob(dit_path + "/**/*", recursive=True)
|
|
|
|
for file_in_d in files:
|
|
if not os.path.isdir(file_in_d):
|
|
content_type = magic.from_file(file_in_d)
|
|
if content_type == "DICOM medical imaging data":
|
|
path = Path(file_in_d)
|
|
with path.open(mode="rb") as f:
|
|
Dicom.objects.create(
|
|
file=File(f, name=file_in_d.split("/")[-1]),
|
|
project=project,
|
|
user=user,
|
|
)
|
|
shutil.rmtree(dit_path)
|
|
tasks.process_project.apply_async(kwargs={"pk": project.pk}, countdown=3)
|
|
return project
|
|
|
|
|
|
def create_coordinate(coordinates, obj):
|
|
for coordinate in coordinates:
|
|
Coordinate.objects.create(
|
|
x=coordinate["x"],
|
|
y=coordinate["y"],
|
|
shape=obj,
|
|
)
|
|
|
|
|
|
def get_bbox(project_id, points, image_range):
|
|
project: Project = Project.objects.get(slug=project_id)
|
|
# print(Dicom.objects.all())
|
|
files = project.files.all()
|
|
bbox_data = []
|
|
for file_number in range(image_range[0], image_range[1] + 1):
|
|
print(points[0]["x"])
|
|
bbox_data.append(
|
|
pydicom.dcmread(files[file_number].file.path)
|
|
.pixel_array[
|
|
int(points[0]["x"]) : int(points[1]["x"]), # noqa
|
|
int(points[0]["y"]) : int(points[1]["y"]), # noqa
|
|
]
|
|
.tolist()
|
|
)
|
|
print(pydicom.dcmread(files[file_number].file.path).pixel_array)
|
|
print(bbox_data)
|
|
# print(project.files.all(), "files", project)
|
|
return []
|
|
|
|
|
|
def generate_3d_point_cloud(project_slug: str):
|
|
project = Project.objects.get(slug=project_slug)
|
|
|
|
point_clouds = []
|
|
|
|
for file_index, file in enumerate(project.files.all()[::10]):
|
|
print(file_index)
|
|
pixel_array = pydicom.dcmread(file.file.path).pixel_array
|
|
for iindex, i in enumerate(pixel_array[::10]):
|
|
for jindex, j in enumerate(i[::10]):
|
|
if j <= 240:
|
|
pass
|
|
# point_clouds.append({'x': jindex, 'y': iindex, 'z': fileindex, 'value': 0})
|
|
else:
|
|
point_clouds.append([jindex, iindex, file_index, j])
|
|
return point_clouds
|
|
|
|
|
|
def generate_3d_model(project: Project, thr=800):
|
|
image = []
|
|
for file in project.files.all():
|
|
image.append(pydicom.dcmread(file.file.path).pixel_array)
|
|
|
|
p = np.array(image).transpose(2, 1, 0)
|
|
verts, faces, normals, values = measure.marching_cubes(p, thr, step_size=1)
|
|
|
|
solid = mesh.Mesh(np.zeros(faces.shape[0], dtype=mesh.Mesh.dtype))
|
|
for i, f in enumerate(faces):
|
|
for j in range(3):
|
|
solid.vectors[i][j] = verts[f[j], :]
|
|
|
|
pth = f"/tmp/{generate_charset(4)}.stl"
|
|
solid.save(pth)
|
|
path = Path(pth)
|
|
with path.open(mode="rb") as f:
|
|
project.stl = File(f, name=pth.split("/")[-1])
|
|
project.save()
|
|
|
|
os.remove(pth)
|