Update app.py

add images to gcs

app.py

@@ -97,8 +97,15 @@ def load_gradcam_model():
 
 # Utility Functions
 
+# Dictionaries to track InstanceNumbers and StudyInstanceUIDs per filename
+# Initialize session state for instance numbers and study UIDs
+if 'instance_numbers' not in st.session_state:
+    st.session_state.instance_numbers = {}
+if 'study_uids' not in st.session_state:
+    st.session_state.study_uids = {}
+
 def upload_to_gcs(image_data: io.BytesIO, filename: str, content_type='application/dicom'):
-    """Uploads an image to Google Cloud Storage."""
+    #Uploads an image to Google Cloud Storage.
     try:
         blob = bucket_result.blob(filename)
         blob.upload_from_file(image_data, content_type=content_type)
@@ -106,9 +113,9 @@ def upload_to_gcs(image_data: io.BytesIO, filename: str, content_type='applicati
     except Exception as e:
         st.error(f"An unexpected error occurred: {e}")
 
-def load_dicom_from_gcs(file_name: str = "dicom_00000001_000.dcm"):
+def load_dicom_from_gcs(dicom_name: str = "dicom_00000001_000.dcm"):
     # Get the blob object
-    blob = bucket_load.blob(file_name)
+    blob = bucket_load.blob(dicom_name)
 
     # Download the file as a bytes object
     dicom_bytes = blob.download_as_bytes()
@@ -116,75 +123,95 @@ def load_dicom_from_gcs(file_name: str = "dicom_00000001_000.dcm"):
     # Wrap bytes object into BytesIO (file-like object)
     dicom_stream = io.BytesIO(dicom_bytes)
 
-    # Load the DICOM file 
+    # Load the DICOM file
     ds = pydicom.dcmread(dicom_stream)
 
     return ds
 
-def png_to_dicom(image_path: str, image_name: str, dicom: str = None):
-    if dicom is None:
-        ds = load_dicom_from_gcs()
-    else:
-        ds = load_dicom_from_gcs(dicom)
+def png_to_dicom(image_path: str, image_name: str, file_name: str, instance_number: int = 1, dicom: str = None, study_instance_uid: str = None, ):
+
+    # Load the template DICOM file
+    ds = load_dicom_from_gcs() if dicom is None else load_dicom_from_gcs(dicom)
 
-    jpg_image = Image.open(image_path)  # Open the image using the path
+    # Process the image
+    jpg_image = Image.open(image_path)  # the PNG or JPG file to be replaced
     print("Image Mode:", jpg_image.mode)
-    if jpg_image.mode == 'L':
-        np_image = np.array(jpg_image.getdata(), dtype=np.uint8)
-        ds.Rows = jpg_image.height
-        ds.Columns = jpg_image.width
-        ds.PhotometricInterpretation = "MONOCHROME1"
-        ds.SamplesPerPixel = 1
-        ds.BitsStored = 8
-        ds.BitsAllocated = 8
-        ds.HighBit = 7
-        ds.PixelRepresentation = 0
-        ds.PixelData = np_image.tobytes()
-        ds.save_as(image_name)
 
-    elif jpg_image.mode == 'RGBA':
-        np_image = np.array(jpg_image.getdata(), dtype=np.uint8)[:, :3]
-        ds.Rows = jpg_image.height
-        ds.Columns = jpg_image.width
-        ds.PhotometricInterpretation = "RGB"
-        ds.SamplesPerPixel = 3
-        ds.BitsStored = 8
-        ds.BitsAllocated = 8
-        ds.HighBit = 7
-        ds.PixelRepresentation = 0
-        ds.PixelData = np_image.tobytes()
-        ds.save_as(image_name)
-    elif jpg_image.mode == 'RGB': 
-        np_image = np.array(jpg_image.getdata(), dtype=np.uint8)[:, :3] # Remove alpha if present
+    if jpg_image.mode in ('L', 'RGBA', 'RGB'):
+        if jpg_image.mode == 'RGBA':
+            np_image = np.array(jpg_image.getdata(), dtype=np.uint8)[:,:3]
+        else:
+            np_image = np.array(jpg_image.getdata(),dtype=np.uint8)
+
         ds.Rows = jpg_image.height
         ds.Columns = jpg_image.width
-        ds.PhotometricInterpretation = "RGB"
-        ds.SamplesPerPixel = 3
+        ds.PhotometricInterpretation = "MONOCHROME1" if jpg_image.mode == 'L' else "RGB"
+        ds.SamplesPerPixel = 1 if jpg_image.mode == 'L' else 3
         ds.BitsStored = 8
         ds.BitsAllocated = 8
         ds.HighBit = 7
         ds.PixelRepresentation = 0
         ds.PixelData = np_image.tobytes()
+
+        if not hasattr(ds, 'PatientName') or ds.PatientName == '':
+            ds.PatientName = os.path.splitext(file_name)[0]  # Remove extension
+
+        ds.SeriesDescription = 'original image' if image_name == 'original_image.dcm' else enhancement_type
+
+        if hasattr(ds, 'StudyDescription'):
+            del ds.StudyDescription
+
+        if study_instance_uid:
+            ds.StudyInstanceUID = study_instance_uid
+        else:
+            # Check if a StudyInstanceUID exists for the file name
+            if file_name in st.session_state.study_uids:
+                ds.StudyInstanceUID = st.session_state.study_uids[file_name]
+                print(f"Reusing StudyInstanceUID for '{file_name}'")
+            else:
+                # Generate a new StudyInstanceUID and store it
+                new_study_uid = generate_uid()
+                st.session_state.study_uids[file_name] = new_study_uid
+                ds.StudyInstanceUID = new_study_uid
+                print(f"New StudyInstanceUID generated for '{file_name}'")
+
+        # Generate a new SeriesInstanceUID and SOPInstanceUID for the added image
+        ds.SeriesInstanceUID = generate_uid()
+        ds.SOPInstanceUID = generate_uid()
+
+        if hasattr(ds, 'InstanceNumber'):
+            st.session_state.instance_numbers[file_name] = int(ds.InstanceNumber) + 1
+        else:
+            # Manage InstanceNumber based on filename
+            if file_name in st.session_state.instance_numbers:
+                st.session_state.instance_numbers[file_name] += 1
+            else:
+                st.session_state.instance_numbers[file_name] = 1
+        ds.InstanceNumber = int(st.session_state.instance_numbers[file_name])
+
         ds.save_as(image_name)
     else:
-        raise ValueError("Unsupported image mode:", jpg_image.mode)
+        raise ValueError(f"Unsupported image mode: {jpg_image.mode}")
+
     return ds
 
-def save_dicom_to_bytes(dicom):
-    dicom_bytes = io.BytesIO()
-    dicom.save_as(dicom_bytes)
-    dicom_bytes.seek(0)
-    return dicom_bytes
+def upload_folder_images(original_image_path, enhanced_image_path, file_name):
+    # Convert images to DICOM if result is png
+    if not original_image_path.lower().endswith('.dcm'):
+        original_dicom = png_to_dicom(original_image_path, "original_image.dcm", file_name=file_name)
+    else:
+        original_dicom = pydicom.dcmread(original_image_path)
+    study_instance_uid = original_dicom.StudyInstanceUID
+
+    # Use StudyInstanceUID as folder name
+    folder_name = study_instance_uid
 
-def upload_folder_images(original_image_path, enhanced_image_path):
-    # Extract the base name of the uploaded image without the extension
-    folder_name = os.path.splitext(uploaded_file.name)[0]
     # Create the folder in Cloud Storage
     bucket_result.blob(folder_name + '/').upload_from_string('', content_type='application/x-www-form-urlencoded')
+
     enhancement_name = enhancement_type.split('_')[-1]
-    # Convert images to DICOM
-    original_dicom = png_to_dicom(original_image_path, "original_image.dcm")
-    enhanced_dicom = png_to_dicom(enhanced_image_path, enhancement_name + ".dcm")
+
+    enhanced_dicom = png_to_dicom(enhanced_image_path, enhancement_name + ".dcm", study_instance_uid=study_instance_uid, file_name=file_name)
 
     # Convert DICOM to byte stream for uploading
     original_dicom_bytes = io.BytesIO()
@@ -199,6 +226,7 @@ def upload_folder_images(original_image_path, enhanced_image_path):
     upload_to_gcs(enhanced_dicom_bytes, folder_name + '/' + enhancement_name + '.dcm', content_type='application/dicom')
 
 
+
 def get_mean_std_per_batch(image_path, H=320, W=320):
     sample_data = []
     for idx, img in enumerate(df.sample(100)["Image Index"].values):
@@ -391,6 +419,7 @@ def redirect_button(url):
 # File uploader for DICOM files
 if uploaded_file is not None:
     if hasattr(uploaded_file, 'name'):
+        file_name = uploaded_file.name
         file_extension = uploaded_file.name.split(".")[-1]  # Get the file extension
         if file_extension.lower() == "dcm":
             # Process DICOM file
@@ -475,7 +504,8 @@ if uploaded_file is not None:
         # Save original image to a file
         original_image_path = "original_image.png"
         cv2.imwrite(original_image_path, original_image)
-
+        upload_folder_images(original_image_path, enhanced_image_path, file_name)
+        
     # Add the redirect button
     col1, col2, col3 = st.columns(3)
     with col1: