1 files changed, 7 insertions, 9 deletions
diff --git a/archaeological_finds/models_treatments.py b/archaeological_finds/models_treatments.py
index f96a0febf..2baced6e1 100644
--- a/archaeological_finds/models_treatments.py
+++ b/archaeological_finds/models_treatments.py
@@ -30,7 +30,8 @@ from django.utils.translation import ugettext_lazy as _, ugettext
 from ishtar_common.utils import cached_label_changed, get_current_year
 from ishtar_common.models import GeneralType, ImageModel, BaseHistorizedItem, \
     OwnPerms, HistoricalRecords, Person, Organization, Source, \
-    ValueGetter, post_save_cache, ShortMenuItem, DashboardFormItem, IshtarImage
+    ValueGetter, post_save_cache, ShortMenuItem, DashboardFormItem, \
+    IshtarImage, ThroughImage
 from archaeological_warehouse.models import Warehouse, Container
 from archaeological_finds.models_finds import Find, FindBasket, TreatmentType
 from archaeological_operations.models import ClosedItem, Operation
@@ -190,8 +191,7 @@ class Treatment(DashboardFormItem, ValueGetter, BaseHistorizedItem,
         return u'{} | {}'.format(u"-".join(items), self.treatment_types_lbl())
 
     def _get_base_image_path(self,):
-        return u"treatment/{}/{}".format(
-            self.year, self.index)
+        return u"{}/{}/{}".format(self.SLUG, self.year, self.index)
 
     def treatment_types_lbl(self):
         """
@@ -307,10 +307,9 @@ def pre_delete_treatment(sender, **kwargs):
 pre_delete.connect(pre_delete_treatment, sender=Treatment)
 
 
-class TreatmentImage(models.Model):
-    image = models.ForeignKey(IshtarImage, on_delete=models.CASCADE)
-    treatment = models.ForeignKey(Treatment, on_delete=models.CASCADE)
-    is_main = models.BooleanField(_(u"Main image"), default=False)
+class TreatmentImage(ThroughImage):
+    item = models.ForeignKey(Treatment, on_delete=models.CASCADE,
+                             related_name='associated_images')
 
 
 class AbsFindTreatments(models.Model):
@@ -642,8 +641,7 @@ class TreatmentSource(Source):
         return self.treatment
 
     def _get_base_image_path(self):
-        return u"treatment/{}/{}/source".format(
-            self.treatment.year, self.treatment.index)
+        return u"{}/source".format(self.treatment._get_base_image_path())
 
 
 class TreatmentFileSource(Source):